% Compute the global solution using a paramerized expectations algorithm

clear
close all
warning off

%% Generate random shocks

t_sim  = 10000;                   % simulation periods
rng('default')
rng(1)                            % same random shock 
shocks = normrnd(0,1,t_sim+1,1);  % random shock
t_drop = 500;                     % number of initial periods to drop in regression

%% Read parameters and steady state

filename = '../linear_model/model_results.mat';
load(filename);
Np = M_.param_nbr;                                            % Number of deep parameters.
for i = 1:Np                                                  % Loop...
  paramname = deblank(M_.param_names(i,:));                   % Get the name of parameter i. 
  eval([ paramname ' = M_.params(' int2str(i) ');']);         % Get the value of parameter i.
end

% 'hU_SS  ' SS labor
% 'LU_SS  ' SS leverage
% 'PU_SS  ' SS run probability
beta    = beta_p;% 'beta_p ' beta
nu      = nu_p;% 'nu_p   ' nu
alpha   = alpha_p;% 'alpha_p' alpha
delta   = delta_p;% 'delta_p' delta
lambda  = lam_p;% 'lam_p  ' lambda
chi0    = chi0_p;% 'chi0_p ' chi0
chi1    = chi1_p;% 'chi1_p ' chi1
rho_a   = rhoa_p;% 'rhoa_p ' rho_a
sigma_a = siga_p;% 'siga_p ' sigma_a

y_ss  = oo_.mean(1);
i_ss  = oo_.mean(2);
xi_ss = oo_.mean(3);
pr_ss = oo_.mean(4);
psi_p = oo_.mean(5);
n0_p  = oo_.mean(6);
gam_p = oo_.mean(7);
k_ss  = oo_.mean(8);
n_ss  = oo_.mean(9);
l_ss  = oo_.mean(10);
rb_ss = oo_.mean(11);
rkhats_ss = oo_.mean(12);
a_ss  = oo_.mean(13);
c_ss  = oo_.mean(14);
h_ss  = oo_.mean(15);

psi   = psi_p;
n0    = n0_p;
gamma = gam_p;

sigma_rk = (1+nu_p)/(1+alpha_p*nu_p)*siga_p;

keep   = lambda*(1-gamma);
lambda = 0.15/0.85;
gamma  = 1-keep/lambda;

lam_p  = lambda;
gam_p  = gamma;

%% Prepare other parameters and polynomial matrix

npf  = 6;  % number of policy functions to approximate
nstv = 4;  % number of state variables
degp = 3;  % degree of polynomials

expmx         = fn_expmx(nstv,degp);    % exponent matrix
expmx_fortran = expmx';

npol = size(expmx,1);

%% initial guess for etas

eta1_rhsee = zeros(npol,1);
eta2_rhsee = zeros(npol,1);
eta3_rhsee = zeros(npol,1);
%eta1_import = dlmread(fullfile('from_matlab_to_fortran/solved_eta1_rhsee.txt'));
%eta2_import = dlmread(fullfile('from_matlab_to_fortran/solved_eta2_rhsee.txt'));
%eta1_import = dlmread(fullfile('from_fortran_to_matlab/temp_eta1_rhsee.txt'));
%eta2_import = dlmread(fullfile('from_fortran_to_matlab/temp_eta2_rhsee.txt'));
%eta3_import = dlmread(fullfile('from_fortran_to_matlab/temp_eta3_rhsee.txt'));
eta1_import = dlmread(fullfile('from_fortran_to_matlab/solved_eta1_rhsee.txt'));
eta2_import = dlmread(fullfile('from_fortran_to_matlab/solved_eta2_rhsee.txt'));
eta3_import = dlmread(fullfile('from_fortran_to_matlab/solved_eta3_rhsee.txt'));
for i=1:npol
    eta1_rhsee(i) = eta1_import(i);
    eta2_rhsee(i) = eta2_import(i);
    eta3_rhsee(i) = eta3_import(i);
end
%eta3_rhsee = eta1_rhsee;

eta1_rbar = zeros(npol,1);
eta2_rbar = zeros(npol,1);
eta3_rbar = zeros(npol,1);
%eta1_import = dlmread(fullfile('from_matlab_to_fortran/solved_eta1_rbar.txt'));
%eta2_import = dlmread(fullfile('from_matlab_to_fortran/solved_eta2_rbar.txt'));
%eta1_import = dlmread(fullfile('from_fortran_to_matlab/temp_eta1_rbar.txt'));
%eta2_import = dlmread(fullfile('from_fortran_to_matlab/temp_eta2_rbar.txt'));
%eta3_import = dlmread(fullfile('from_fortran_to_matlab/temp_eta3_rbar.txt'));
eta1_import = dlmread(fullfile('from_fortran_to_matlab/solved_eta1_rbar.txt'));
eta2_import = dlmread(fullfile('from_fortran_to_matlab/solved_eta2_rbar.txt'));
eta3_import = dlmread(fullfile('from_fortran_to_matlab/solved_eta3_rbar.txt'));

for i=1:npol
    eta1_rbar(i) = eta1_import(i);
    eta2_rbar(i) = eta2_import(i);
    eta3_rbar(i) = eta3_import(i);
end
%eta3_rbar = eta1_rbar;

%% Set parameters for numerical integration and iterations

nd = 70;  % nd percent drop in net worth in a crisis

n_trapz = 301;
z_max   = 6;
z_min   = -z_max;
z_trapz = linspace(z_min,z_max,n_trapz);
z_trapz = z_trapz';

iter_max = 200;
tol      = 2e-4;   % convergence criterion
adj      = 0.30;   % adjustment speed of polynomial coefficients, the higher the larger update and the faster.

%% Prepare for welfare analysis

length = 3000;
repeat = 1000;
welfare_integer = [length,repeat];
save('from_matlab_to_fortran/welfare_integer.txt','welfare_integer','-ascii','-double');

n1   = n_ss;
k1   = k_ss;
rb1  = rb_ss;
a1   = 0;
welfare_states = [n1,k1,rb1,a1];
save('from_matlab_to_fortran/welfare_states.txt','welfare_states','-ascii','-double');

%% Prepare for fortran

parameters    = [hU_SS,LU_SS,PU_SS,beta_p,nu_p,alpha_p,delta_p,lam_p,chi0_p,chi1_p,rhoa_p,siga_p,nd];
save('from_matlab_to_fortran/parameters.txt','parameters','-ascii');

parameters_ss = [c_ss,y_ss,i_ss,k_ss,h_ss,n_ss,l_ss,rb_ss,rkhats_ss,a_ss,xi_ss,pr_ss,psi_p,n0_p,gam_p,sigma_rk];
save('from_matlab_to_fortran/ss_values.txt','parameters_ss','-ascii');

parameters_integer = [npf,nstv,degp,npol,iter_max,n_trapz,t_sim,t_drop];
save('from_matlab_to_fortran/parameters_integer.txt','parameters_integer','-ascii');

parameters_solution = [tol,adj];
save('from_matlab_to_fortran/parameters_solution.txt','parameters_solution','-ascii');

save('from_matlab_to_fortran/polynomial_matrix.txt','expmx_fortran','-ascii','-double');
save('from_matlab_to_fortran/z_trapz.txt','z_trapz','-ascii','-double');
save('from_matlab_to_fortran/shocks.txt','shocks','-ascii','-double');

%% run fortran

skip = 0; % 1 implies skipping fortran simulation and just import 'temp' etas.

if skip ~= 1

for i=1:iter_max
    
    % export etas to fortran
    save('from_matlab_to_fortran/eta1_rhsee.txt','eta1_rhsee','-ascii','-double');
    save('from_matlab_to_fortran/eta2_rhsee.txt','eta2_rhsee','-ascii','-double');
    save('from_matlab_to_fortran/eta3_rhsee.txt','eta3_rhsee','-ascii','-double');
    save('from_matlab_to_fortran/eta1_rbar.txt','eta1_rbar','-ascii','-double');
    save('from_matlab_to_fortran/eta2_rbar.txt','eta2_rbar','-ascii','-double');
    save('from_matlab_to_fortran/eta3_rbar.txt','eta3_rbar','-ascii','-double');
    
    % simulate using fortran
%    command='/home/matlab/Desktop/Matsumoto/Bank_Run/20200529_policy_unix/matfort_nd70_2nd/main.exe';
%    command='./main.exe';
%    [status,cmdout] = unix(command);
    
    system('x64\Release\f90_bankrun_pea.exe');

    % import rhsee and rbar derived, and sequence of state variables x
    rhsee_norun = dlmread(fullfile('from_fortran_to_matlab/rhsee_norun.txt'));
    rhsee_run   = dlmread(fullfile('from_fortran_to_matlab/rhsee_run.txt'));
    rhsee_nega  = dlmread(fullfile('from_fortran_to_matlab/rhsee_nega.txt'));
    rbar_norun  = dlmread(fullfile('from_fortran_to_matlab/rbar_norun.txt'));
    rbar_run    = dlmread(fullfile('from_fortran_to_matlab/rbar_run.txt'));
    rbar_nega   = dlmread(fullfile('from_fortran_to_matlab/rbar_nega.txt'));
    x_norun       = dlmread(fullfile('from_fortran_to_matlab/x_norun.txt'));
        num_norun = size(x_norun,1)/npol;
        x_norun   = reshape(x_norun,num_norun,npol);
    x_run         = dlmread(fullfile('from_fortran_to_matlab/x_run.txt'));
        num_run   = size(x_run,1)/npol;
        x_run     = reshape(x_run,num_run,npol);
    x_nega        = dlmread(fullfile('from_fortran_to_matlab/x_nega.txt'));
        num_nega  = size(x_nega,1)/npol;
        x_nega    = reshape(x_nega,num_nega,npol);
    
    % regress    
    eta1_rhsee_new = (x_norun'*x_norun)\x_norun'*rhsee_norun;
    eta2_rhsee_new = (x_run'*  x_run  )\x_run'*  rhsee_run;
    eta3_rhsee_new = (x_nega'* x_nega )\x_nega'* rhsee_nega;
    eta1_rbar_new  = (x_norun'*x_norun)\x_norun'*rbar_norun;
    eta2_rbar_new  = (x_run'*  x_run  )\x_run'*  rbar_run;
    eta3_rbar_new  = (x_nega'* x_nega )\x_nega'* rbar_nega;
    
    % check convergence
    
    con_lev = zeros(6,1);
    rsq     = zeros(6,1);
    
    con_lev(1) = max( abs( (exp(x_norun*eta1_rhsee_new) - exp(x_norun*eta1_rhsee)) ./ exp(x_norun*eta1_rhsee) ) );
    con_lev(2) = max( abs( (exp(x_run  *eta2_rhsee_new) - exp(x_run  *eta2_rhsee)) ./ exp(x_run  *eta2_rhsee) ) );
    con_lev(3) = max( abs( (exp(x_nega *eta3_rhsee_new) - exp(x_nega *eta3_rhsee)) ./ exp(x_nega *eta3_rhsee) ) );
    con_lev(4) = max( abs( (exp(x_norun*eta1_rbar_new ) - exp(x_norun*eta1_rbar )) ./ exp(x_norun*eta1_rbar ) ) );
    con_lev(5) = max( abs( (exp(x_run  *eta2_rbar_new ) - exp(x_run  *eta2_rbar )) ./ exp(x_run  *eta2_rbar ) ) );
    con_lev(6) = max( abs( (exp(x_nega *eta3_rbar_new ) - exp(x_nega *eta3_rbar )) ./ exp(x_nega *eta3_rbar) ) );

    rsq(1) = 1 - sum((rhsee_norun - x_norun*eta1_rhsee).^2) ./ sum((rhsee_norun-mean(rhsee_norun)).^2);
    rsq(2) = 1 - sum((rhsee_run   - x_run  *eta2_rhsee).^2) ./ sum((rhsee_run  -mean(rhsee_run  )).^2);
    rsq(3) = 1 - sum((rhsee_nega  - x_nega *eta3_rhsee).^2) ./ sum((rhsee_nega -mean(rhsee_nega )).^2);
    rsq(4) = 1 - sum((rbar_norun  - x_norun*eta1_rbar ).^2) ./ sum((rbar_norun -mean(rbar_norun )).^2);
    rsq(5) = 1 - sum((rbar_run    - x_run  *eta2_rbar ).^2) ./ sum((rbar_run   -mean(rbar_run   )).^2);
    rsq(6) = 1 - sum((rbar_nega   - x_nega *eta3_rbar ).^2) ./ sum((rbar_nega  -mean(rbar_nega  )).^2);

    % display results
    fprintf('#################### iteration:')
    disp(i)

    fprintf('max gap in convergence:')
    disp(max(con_lev(1:3)))
    
    fprintf('R-square:')
    disp(rsq')
    
    % temp solution is saved in fortran code

    if max(con_lev(1:3)) < tol
        break
    end
    
    % update etas
    eta1_rhsee = adj*eta1_rhsee_new + (1-adj)*eta1_rhsee;
    eta2_rhsee = adj*eta2_rhsee_new + (1-adj)*eta2_rhsee;
    eta3_rhsee = adj*eta3_rhsee_new + (1-adj)*eta3_rhsee;
    eta1_rbar  = adj*eta1_rbar_new  + (1-adj)*eta1_rbar;
    eta2_rbar  = adj*eta2_rbar_new  + (1-adj)*eta2_rbar;
    eta3_rbar  = adj*eta3_rbar_new  + (1-adj)*eta3_rbar;

end

end

return
